package utils

import (
	"sync"
	"sync/atomic"
)

// WaitGroup 用来管理一组同时以并发或串行执行的函数
//   - 如果以串行执行，不需要调用 [WaitGroup.Done]
//   - 如果以并发执行，需要自行调用 [WaitGroup.Done]
type WaitGroup struct {
	lock   sync.Mutex
	wg     sync.WaitGroup
	i      int32
	Thread bool
}

var wgpool = sync.Pool{
	New: func() interface{} {
		return &WaitGroup{}
	},
}

// NewWaitGroup 从 [sync.Pool] 中获取一个 [*WaitGroup]
func NewWaitGroup(Thread bool) *WaitGroup {
	ret := wgpool.Get().(*WaitGroup)
	ret.Thread = Thread
	return ret
}

// Add 记录新增执行一个函数
func (wg *WaitGroup) Add() {
	if wg.Thread {
		wg.wg.Add(1)
		return
	}
	wg.lock.Lock()
	atomic.AddInt32(&wg.i, 1)
}

// Done 记录一个函数执行完毕
func (wg *WaitGroup) Done() {
	if wg.Thread {
		wg.wg.Done()
		return
	}
	if atomic.LoadInt32(&wg.i) != 0 {
		wg.lock.Unlock()
	}
}

// Wait 等待所有函数执行完毕
func (wg *WaitGroup) Wait() {
	if wg.Thread {
		wg.wg.Wait()
		return
	}
	wg.lock.Lock()
}

// GoI 执行一个函数，该函数接受一个int型参数，该函数无需调用 [WaitGroup.Add]
//   - 如果以串行执行，调用传入函数，执行完毕后返回
//   - 如果以并发执行，安排传入函数在一个goroutine执行，然后直接返回
func (wg *WaitGroup) GoI(f func(gi int), i int) {
	if wg.Thread {
		wg.wg.Add(1)
		go f(i)
		return
	}
	f(i)
}

// GoI2 执行一个函数，该函数接受两个int型参数，该函数无需调用 [WaitGroup.Add]
//   - 如果以串行执行，调用传入函数，执行完毕后返回
//   - 如果以并发执行，安排传入函数在一个goroutine执行，然后直接返回
func (wg *WaitGroup) GoI2(f func(i1, i2 int), i1, i2 int) {
	if wg.Thread {
		wg.wg.Add(1)
		go f(i1, i2)
		return
	}
	f(i1, i2)
}

// Go 执行一个函数，该函数没有参数，该函数无需调用 [WaitGroup.Add]
//   - 如果以串行执行，调用传入函数，执行完毕后返回
//   - 如果以并发执行，安排传入函数在一个goroutine执行，然后直接返回
func (wg *WaitGroup) Go(f func()) {
	if wg.Thread {
		wg.wg.Add(1)
		go f()
		return
	}
	f()
}

// Close 将 一个 [*WaitGroup] 放入 [sync.Pool] 中
func (wg *WaitGroup) Close() {
	//可能修改的恢复零值
	wg.i = 0
	wg.lock = sync.Mutex{}
	wgpool.Put(wg) //放入池中
}
